import cv2
import os
import numpy as np

def getImagesAndLabels(path):
    faceSamples = []
    ids = []

    # 遍历faces文件夹中的所有子文件夹
    for root, dirs, _ in os.walk(path):
        for dir_name in dirs:
            try:
                # 尝试将子文件夹名称转换为整数
                id = int(dir_name)
                # 获取子文件夹中的所有图像文件
                for img_file in os.listdir(os.path.join(root, dir_name)):
                    PIL_img = cv2.imread(os.path.join(root, dir_name, img_file), cv2.IMREAD_GRAYSCALE)
                    if PIL_img is not None:  # 确保图像成功加载
                        faces = cascade.detectMultiScale(PIL_img, scaleFactor=1.3, minNeighbors=5)

                        for (x, y, w, h) in faces:
                            faceSamples.append(PIL_img[y:y + h, x:x + w])
                            ids.append(id)
            except ValueError:
                # 如果子文件夹名称不是数字，则忽略
                continue

    return faceSamples, ids

# 设置训练数据路径
train_path = 'faces'

# 创建Haar级联分类器对象
cascade = cv2.CascadeClassifier('data/haarcascade_frontalface_alt.xml')

# 获取图像和标签
faces, ids = getImagesAndLabels(train_path)
ids = np.array(ids)

# 创建LBPHFaceRecognizer对象并训练模型
recognizer = cv2.face.LBPHFaceRecognizer_create()
recognizer.train(faces, ids)

# 保存模型
recognizer.save('trainer.yml')
